package main

/*
 *  Go kprobe examale.
 */

import (
	"fmt"
	"io/ioutil"
	"os"
	"strconv"
	"strings"
	"time"

	"golang.org/x/sys/unix"

	"github.com/cilium/ebpf"
)

//go:generate go run github.com/cilium/ebpf/cmd/bpf2go pinger pinger.c -- -I./include  -nostdinc -O3

// Header used to compile examples
//-I.linux_kernel/arch/x86/include -I./linux_kernel/arch/x86/include/generated  -I./linux_kernel/include -I./linux_kernel/arch/x86/include/uapi -I./linux_kernel/arch/x86/include/generated/uapi -I./linux_kernel//include/uapi -I./linux_kernel/include/generated/uapi

const trace_point_name = "kprobe_icmp_echo"

func main() {
	err := unix.Setrlimit(unix.RLIMIT_MEMLOCK, &unix.Rlimit{
		Cur: unix.RLIM_INFINITY,
		Max: unix.RLIM_INFINITY,
	})
	if err != nil {
		fmt.Println("WARNING: Failed to adjust rlimit")
	}

	specs, err := newPingerSpecs()
	if err != nil {
		panic(err)
	}

	objs, err := specs.Load(nil)
	if err != nil {
		panic(err)
	}
	defer objs.Close()

	// Set up the map
	events, err := ebpf.NewMap(specs.MapPingsMap)

	if err != nil {
		panic(fmt.Errorf("could not create event map: %v\n", err))
	}
	defer events.Close()

	prog := objs.ProgramDumpPings
	if err != nil {
		panic(fmt.Errorf("could not create new ebpf program: %v", err))
	}
	defer prog.Close()

	/* Set up the trace point so we can get an ID for it */
	err = createTracePoint(trace_point_name, "icmp_echo")
	if err != nil {
		panic(fmt.Errorf("could not create new kprobe trace point: %v", err))
	}

	defer func() {
		removeTracePoint(trace_point_name)
	}()

	err = enableTracePoint(trace_point_name)
	if err != nil {
		panic(fmt.Errorf("could not enable kprobe trace point: %v", err))
	}
	defer func() {
		disableTracePoint(trace_point_name)
	}()

	tid, err := getTracepointID(trace_point_name)
	if err != nil {
		panic(fmt.Errorf("could not get tracepoint id: %v", err))
	}
	fmt.Printf("Trace point id: %d\n", tid)

	attr := unix.PerfEventAttr{
		Type:        unix.PERF_TYPE_TRACEPOINT,
		Config:      tid,
		Sample_type: unix.PERF_SAMPLE_RAW,
		Sample:      1,
		Wakeup:      1,
	}
	pfd, err := unix.PerfEventOpen(&attr, -1, 0, -1, unix.PERF_FLAG_FD_CLOEXEC)
	if err != nil {
		panic(fmt.Errorf("unable to open perf events: %v", err))
	}
	if _, _, errno := unix.Syscall(unix.SYS_IOCTL, uintptr(pfd), unix.PERF_EVENT_IOC_ENABLE, 0); errno != 0 {
		panic(fmt.Errorf("unable to enable perf events: %v", err))
	}
	if _, _, errno := unix.Syscall(unix.SYS_IOCTL, uintptr(pfd), unix.PERF_EVENT_IOC_SET_BPF, uintptr(prog.FD())); errno != 0 {
		panic(fmt.Errorf("unable to attach bpf program to perf events: %v", err))
	}

	// Every few seconds dump the map output
	for range time.Tick(3 * time.Second) {
		iter := objs.MapPingsMap.Iterate()
		var key uint32
		var values uint32
		for iter.Next(&key, &values) {
			fmt.Printf("key: %d      count: %d\n", key, values)
		}
	}

	if _, _, errno := unix.Syscall(unix.SYS_IOCTL, uintptr(pfd), unix.PERF_EVENT_IOC_DISABLE, 0); errno != 0 {
		panic(fmt.Errorf("unable to disable perf events: %v", err))
	}
}

// getTracepointID returns the system specific ID for the tracepoint sched_process_exec
func getTracepointID(name string) (uint64, error) {
	fname := fmt.Sprintf("/sys/kernel/tracing/events/kprobes/%s/id", name)
	data, err := ioutil.ReadFile(fname)
	// /sys/kernel/debug/tracing/events/sched/sched_process_exec/id")
	if err != nil {
		return 0, fmt.Errorf("failed to read tracepoint ID for '%s': %v", name, err)
	}
	tid := strings.TrimSuffix(string(data), "\n")
	return strconv.ParseUint(tid, 10, 64)
}

// Need to debug this and make sure we're setting up the arugments correctly.
//echo 'p:isn_sock_probe sock_sendmsg sock=%di msg=%si'   >> /sys/kernel/debug/tracing/kprobe_events
// echo 'p:isn_icmp_echo icmp_echo'   >> /sys/kernel/debug/tracing/kprobe_events
func createTracePoint(name string, kprobefunction string) error {
	msg := fmt.Sprintf("p:%s %s  skb=%%di", name, kprobefunction)
	return write_kprobe_events(msg)
}

func removeTracePoint(name string) error {
	msg := fmt.Sprintf("-:%s", name)
	return write_kprobe_events(msg)
}

func write_kprobe_events(msg string) error {
	f, err := os.OpenFile("/sys/kernel/debug/tracing/kprobe_events", os.O_WRONLY, 0644)
	if err != nil {
		return err
	}
	defer f.Close()
	f.WriteString(msg)

	return nil
}

func toggleTracePoint(name string, msg string) error {
	enablefile := fmt.Sprintf("/sys/kernel/debug/tracing/events/kprobes/%s/enable", name)
	f, err := os.OpenFile(enablefile, os.O_WRONLY, 0644)
	if err != nil {
		return err
	}
	defer f.Close()
	f.WriteString(msg)
	return nil
}

func disableTracePoint(name string) error {
	return toggleTracePoint(name, "0")
}

func enableTracePoint(name string) error {
	return toggleTracePoint(name, "1")
}
